package utils

import (
	"gorm.io/gorm"
	"log"
	"os"
	"path/filepath"
	"sync"

	"goscan-mgt/config"

	"github.com/glebarez/sqlite"
)

var (
	DB   *gorm.DB
	once sync.Once
)

func GetDB() *gorm.DB {
	var err error
	once.Do(func() {
		cfg := config.LoadConfig()
		switch cfg.DB.Type {
		case "sqlite":
			dir := filepath.Dir(cfg.DB.Path)
			if err := os.MkdirAll(dir, 0755); err != nil {
				log.Fatalf("创建数据库目录失败: %v", err)
			}
			DB, err = gorm.Open(sqlite.Open(cfg.DB.Path), &gorm.Config{})
		default:
			log.Fatalf("不支持的数据库类型: %s", cfg.DB.Type)
		}
		if err != nil {
			log.Fatalf("数据库连接失败: %v", err)
		}
	})
	return DB
}

// 初始化数据库
func InitDB() {
	db := GetDB()
	sqlDB, err := db.DB() // 获取底层 *sql.DB
	if err != nil {
		log.Fatalf("无法获取底层数据库对象: %v", err)
	}
	if err = sqlDB.Ping(); err != nil {
		log.Fatalf("无法访问数据库: %v", err)
	}
	initTables()
}

// 初始化数据表
func initTables() {
	// 创建服务信息表
	createServiceInfoTable := `
    CREATE TABLE IF NOT EXISTS service_infos (
        id INTEGER PRIMARY KEY,
        port INTEGER NOT NULL,
        service VARCHAR(100) NOT NULL,
        create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		UNIQUE (port, service)  -- 添加唯一约束，确保组合不重复
    );`
	if err := DB.Exec(createServiceInfoTable).Error; err != nil {
		log.Fatalf("创建服务信息表失败: %v", err)
	}

	// 创建项目表
	createProjectTable := `
    CREATE TABLE IF NOT EXISTS projects (
        id INTEGER PRIMARY KEY,
        project_name TEXT NOT NULL UNIQUE,
        scan_count INTEGER NOT NULL DEFAULT 0, --扫描条数
        create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        status INTEGER NOT NULL DEFAULT 0 CHECK (status IN (0, 1)) -- 0: 扫描中, 1: 已完成
    );`
	if err := DB.Exec(createProjectTable).Error; err != nil {
		log.Fatalf("创建项目表失败: %v", err)
	}

	// 创建扫描结果表
	createScanResultTable := `
    CREATE TABLE IF NOT EXISTS scan_results (
        id INTEGER PRIMARY KEY,
        ip TEXT NOT NULL,
        port INTEGER NOT NULL,
        port_status VARCHAR(20) NOT NULL,
        service VARCHAR(100),
        create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        project_id INTEGER,
        FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
    );`
	if err := DB.Exec(createScanResultTable).Error; err != nil {
		log.Fatalf("创建扫描结果表失败: %v", err)
	}
	// 初始化常见服务与端口信息
	initServiceInfo()
}

// 初始化常见服务与端口信息
func initServiceInfo() {
	var count int64
	if err := DB.Table("service_infos").Count(&count).Error; err != nil {
		log.Printf("查询服务信息表失败: %v", err)
		return
	}
	if count > 0 {
		log.Printf("service_infos 表中已经存在: %v 条数据。", count)
		return
	}

	insertServiceInfo := `
    INSERT INTO service_infos (port, service) VALUES
    (80, 'HTTP'),
    (443, 'HTTPS'),
    (21, 'FTP'),
    (22, 'SSH'),
    (23, 'Telnet'),
    (25, 'SMTP'),
    (53, 'DNS'),
    (110, 'POP3'),
    (143, 'IMAP'),
    (3306, 'MySQL'),
    (5432, 'PostgreSQL'),
    (6379, 'Redis'),
    (9200, 'Elasticsearch'),
    (5601, 'Kibana'),
    (5672, 'RabbitMQ'),
    (15672, 'RabbitMQ-Management'),
    (9092, 'Kafka'),
    (27017, 'MongoDB'),
    (8080, 'HTTP-Proxy'),
    (8081, 'HTTP-Alternate'),
    (9000, 'Nginx-HTTP'),
    (9001, 'Nginx-Admin'),
    (8000, 'Web Server'),
    (3000, 'Node.js App'),
    (5000, 'Flask App'),
    (8089, 'Solr'),
    (8765, 'ActiveMQ'),
    (11211, 'Memcached'),
    (50070, 'Hadoop NameNode'),
    (8088, 'Hadoop ResourceManager'),
    (9042, 'Cassandra'),
    (9090, 'Prometheus'),
    (3001, 'Grafana'),
    (8500, 'Consul'),
    (8086, 'InfluxDB'),
    (9093, 'Kafka Connect'),
    (2181, 'ZooKeeper');`

	if err := DB.Exec(insertServiceInfo).Error; err != nil {
		log.Printf("插入服务信息失败: %v", err)
	}
}
